{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Notebook Setup \n",
    "The following cell will install Drake, checkout the underactuated repository, and set up the path (only if necessary).\n",
    "- On Google's Colaboratory, this **will take approximately two minutes** on the first time it runs (to provision the machine), but should only need to reinstall once every 12 hours.  Colab will ask you to \"Reset all runtimes\"; say no to save yourself the reinstall.\n",
    "- On Binder, the machines should already be provisioned by the time you can run this; it should return (almost) instantly.\n",
    "\n",
    "More details are available [here](http://underactuated.mit.edu/drake.html)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "try:\n",
    "  import pydrake\n",
    "  import underactuated\n",
    "except ImportError:\n",
    "  !curl -s https://raw.githubusercontent.com/RussTedrake/underactuated/master/scripts/setup/jupyter_setup.py > jupyter_setup.py\n",
    "  from jupyter_setup import setup_underactuated\n",
    "  setup_underactuated()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## LQR for the Acrobot"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "\n",
    "from pydrake.all import (DiagramBuilder, LinearQuadraticRegulator,\n",
    "                         PlanarSceneGraphVisualizer, Saturation,\n",
    "                         SceneGraph, Simulator, WrapToSystem)\n",
    "from pydrake.examples.acrobot import (AcrobotInput, AcrobotGeometry, \n",
    "                                      AcrobotPlant, AcrobotState)\n",
    "from underactuated.jupyter import AdvanceToAndVisualize, SetupMatplotlibBackend, running_as_notebook\n",
    "plt_is_interactive = SetupMatplotlibBackend()\n",
    "\n",
    "def UprightState():\n",
    "    state = AcrobotState()\n",
    "    state.set_theta1(np.pi)\n",
    "    state.set_theta2(0.)\n",
    "    state.set_theta1dot(0.)\n",
    "    state.set_theta2dot(0.)\n",
    "    return state\n",
    "\n",
    "\n",
    "def BalancingLQR():\n",
    "    # Design an LQR controller for stabilizing the Acrobot around the upright.\n",
    "    # Returns a (static) AffineSystem that implements the controller (in\n",
    "    # the original AcrobotState coordinates).\n",
    "\n",
    "    acrobot = AcrobotPlant()\n",
    "    context = acrobot.CreateDefaultContext()\n",
    "\n",
    "    input = AcrobotInput()\n",
    "    input.set_tau(0.)\n",
    "    context.FixInputPort(0, input)\n",
    "\n",
    "    context.get_mutable_continuous_state_vector()\\\n",
    "        .SetFromVector(UprightState().CopyToVector())\n",
    "\n",
    "    Q = np.diag((10., 10., 1., 1.))\n",
    "    R = [1]\n",
    "\n",
    "    return LinearQuadraticRegulator(acrobot, context, Q, R)\n",
    "\n",
    "\n",
    "builder = DiagramBuilder()\n",
    "acrobot = builder.AddSystem(AcrobotPlant())\n",
    "\n",
    "saturation = builder.AddSystem(Saturation(min_value=[-10], max_value=[10]))\n",
    "builder.Connect(saturation.get_output_port(0), acrobot.get_input_port(0))\n",
    "wrapangles = WrapToSystem(4)\n",
    "wrapangles.set_interval(0, 0, 2. * np.pi)\n",
    "wrapangles.set_interval(1, -np.pi, np.pi)\n",
    "wrapto = builder.AddSystem(wrapangles)\n",
    "builder.Connect(acrobot.get_output_port(0), wrapto.get_input_port(0))\n",
    "controller = builder.AddSystem(BalancingLQR())\n",
    "builder.Connect(wrapto.get_output_port(0), controller.get_input_port(0))\n",
    "builder.Connect(controller.get_output_port(0), saturation.get_input_port(0))\n",
    "\n",
    "# Setup visualization\n",
    "scene_graph = builder.AddSystem(SceneGraph())\n",
    "AcrobotGeometry.AddToBuilder(builder, acrobot.get_output_port(0), scene_graph)\n",
    "visualizer = builder.AddSystem(\n",
    "    PlanarSceneGraphVisualizer(scene_graph, xlim=[-4., 4.], ylim=[-4., 4.], show=plt_is_interactive))\n",
    "builder.Connect(scene_graph.get_pose_bundle_output_port(),\n",
    "                visualizer.get_input_port(0))\n",
    "\n",
    "diagram = builder.Build()\n",
    "\n",
    "# Set up a simulator to run this diagram\n",
    "simulator = Simulator(diagram)\n",
    "context = simulator.get_mutable_context()\n",
    "\n",
    "# Simulate\n",
    "duration = 4.0 if running_as_notebook else 0.1 # sets a shorter duration during testing\n",
    "if not visualizer._show:\n",
    "    visualizer.start_recording()\n",
    "for i in range(5):\n",
    "    context.SetTime(0.)\n",
    "    context.SetContinuousState(UprightState().CopyToVector() +\n",
    "                               0.05 * np.random.randn(4,))\n",
    "    simulator.Initialize()\n",
    "    if i<4:\n",
    "        simulator.AdvanceTo(duration)\n",
    "    else:\n",
    "        AdvanceToAndVisualize(simulator, visualizer, duration)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## LQR for the Cart-Pole"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "\n",
    "from pydrake.all import (AddMultibodyPlantSceneGraph, DiagramBuilder, \n",
    "                         LinearQuadraticRegulator, Parser, PlanarSceneGraphVisualizer, \n",
    "                         SceneGraph, Simulator)\n",
    "from underactuated import FindResource\n",
    "from underactuated.jupyter import AdvanceToAndVisualize, SetupMatplotlibBackend, running_as_notebook\n",
    "plt_is_interactive = SetupMatplotlibBackend()\n",
    "\n",
    "def UprightState():\n",
    "    state = (0, np.pi, 0, 0)\n",
    "    return state\n",
    "\n",
    "\n",
    "def BalancingLQR(plant):\n",
    "    # Design an LQR controller for stabilizing the CartPole around the upright.\n",
    "    # Returns a (static) AffineSystem that implements the controller (in\n",
    "    # the original CartPole coordinates).\n",
    "\n",
    "    context = plant.CreateDefaultContext()\n",
    "    plant.get_actuation_input_port().FixValue(context, [0])\n",
    "\n",
    "    context.get_mutable_continuous_state_vector().SetFromVector(UprightState())\n",
    "\n",
    "    Q = np.diag((10., 10., 1., 1.))\n",
    "    R = [1]\n",
    "\n",
    "    # MultibodyPlant has many (optional) input ports, so we must pass the\n",
    "    # input_port_index to LQR.\n",
    "    return LinearQuadraticRegulator(\n",
    "        plant,\n",
    "        context,\n",
    "        Q,\n",
    "        R,\n",
    "        input_port_index=plant.get_actuation_input_port().get_index())\n",
    "\n",
    "\n",
    "builder = DiagramBuilder()\n",
    "plant, scene_graph = AddMultibodyPlantSceneGraph(builder, time_step=0.0)\n",
    "file_name = FindResource(\"models/cartpole.urdf\")\n",
    "Parser(plant).AddModelFromFile(file_name)\n",
    "plant.Finalize()\n",
    "\n",
    "controller = builder.AddSystem(BalancingLQR(plant))\n",
    "builder.Connect(plant.get_state_output_port(), controller.get_input_port(0))\n",
    "builder.Connect(controller.get_output_port(0),\n",
    "                plant.get_actuation_input_port())\n",
    "\n",
    "# Setup visualization\n",
    "visualizer = builder.AddSystem(\n",
    "    PlanarSceneGraphVisualizer(scene_graph, xlim=[-2.5, 2.5], ylim=[-1, 2.5], show=plt_is_interactive))\n",
    "builder.Connect(scene_graph.get_pose_bundle_output_port(),\n",
    "                visualizer.get_input_port(0))\n",
    "\n",
    "diagram = builder.Build()\n",
    "\n",
    "# Set up a simulator to run this diagram\n",
    "simulator = Simulator(diagram)\n",
    "context = simulator.get_mutable_context()\n",
    "\n",
    "# Simulate\n",
    "duration = 8.0 if running_as_notebook else 0.1 # sets a shorter duration during testing\n",
    "if not visualizer._show:\n",
    "    visualizer.start_recording()\n",
    "for i in range(5):\n",
    "    context.SetTime(0.)\n",
    "    context.SetContinuousState(UprightState() + 0.1 * np.random.randn(4,))\n",
    "    simulator.Initialize()\n",
    "    if i<4:\n",
    "        simulator.AdvanceTo(duration)\n",
    "    else:\n",
    "        AdvanceToAndVisualize(simulator, visualizer, duration)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## LQR for the Planar Quadrotor"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "\n",
    "from pydrake.all import DiagramBuilder, LinearQuadraticRegulator, Simulator\n",
    "from underactuated.quadrotor2d import Quadrotor2D, Quadrotor2DVisualizer\n",
    "from underactuated.jupyter import AdvanceToAndVisualize, SetupMatplotlibBackend, running_as_notebook\n",
    "plt_is_interactive = SetupMatplotlibBackend()\n",
    "\n",
    "\n",
    "def QuadrotorLQR(plant):\n",
    "    context = plant.CreateDefaultContext()\n",
    "    context.SetContinuousState(np.zeros([6, 1]))\n",
    "    context.FixInputPort(0, plant.mass * plant.gravity / 2. * np.array([1, 1]))\n",
    "\n",
    "    Q = np.diag([10, 10, 10, 1, 1, (plant.length / 2. / np.pi)])\n",
    "    R = np.array([[0.1, 0.05], [0.05, 0.1]])\n",
    "\n",
    "    return LinearQuadraticRegulator(plant, context, Q, R)\n",
    "\n",
    "\n",
    "builder = DiagramBuilder()\n",
    "plant = builder.AddSystem(Quadrotor2D())\n",
    "\n",
    "controller = builder.AddSystem(QuadrotorLQR(plant))\n",
    "builder.Connect(controller.get_output_port(0), plant.get_input_port(0))\n",
    "builder.Connect(plant.get_output_port(0), controller.get_input_port(0))\n",
    "\n",
    "# Setup visualization\n",
    "visualizer = builder.AddSystem(Quadrotor2DVisualizer(show=plt_is_interactive))\n",
    "builder.Connect(plant.get_output_port(0), visualizer.get_input_port(0))\n",
    "\n",
    "diagram = builder.Build()\n",
    "\n",
    "# Set up a simulator to run this diagram\n",
    "simulator = Simulator(diagram)\n",
    "context = simulator.get_mutable_context()\n",
    "\n",
    "# Simulate\n",
    "duration = 4.0 if running_as_notebook else 0.1 # sets a shorter duration during testing\n",
    "if not visualizer._show:\n",
    "    visualizer.start_recording()\n",
    "for i in range(5):\n",
    "    context.SetTime(0.)\n",
    "    context.SetContinuousState(np.random.randn(6,))\n",
    "    simulator.Initialize()\n",
    "    if i<4:\n",
    "        simulator.AdvanceTo(duration)\n",
    "    else:\n",
    "        AdvanceToAndVisualize(simulator, visualizer, duration)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## LQR for the 3D Quadrotor\n",
    "\n",
    "This example requires a 3D visualizer, and does not work [yet](https://github.com/rdeits/meshcat/issues/67#issuecomment-632453885) on Colab (but should work on Binder).  "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "\n",
    "from pydrake.all import (DiagramBuilder,\n",
    "                         MeshcatVisualizer,\n",
    "                         SceneGraph,\n",
    "                         Simulator,\n",
    "                         VectorSystem)\n",
    "from pydrake.examples.quadrotor import (QuadrotorPlant,\n",
    "                                        StabilizingLQRController,\n",
    "                                        QuadrotorGeometry)\n",
    "from underactuated.jupyter import running_as_notebook\n",
    "\n",
    "builder = DiagramBuilder()\n",
    "\n",
    "plant = builder.AddSystem(QuadrotorPlant())\n",
    "\n",
    "controller = builder.AddSystem(StabilizingLQRController(plant, [0, 0, 1]))\n",
    "builder.Connect(controller.get_output_port(0), plant.get_input_port(0))\n",
    "builder.Connect(plant.get_output_port(0), controller.get_input_port(0))\n",
    "\n",
    "# Set up visualization in MeshCat\n",
    "scene_graph = builder.AddSystem(SceneGraph())\n",
    "QuadrotorGeometry.AddToBuilder(builder, plant.get_output_port(0), scene_graph)\n",
    "visualizer = builder.AddSystem(\n",
    "    MeshcatVisualizer(scene_graph, zmq_url=\"new\", open_browser=False))\n",
    "builder.Connect(scene_graph.get_pose_bundle_output_port(),\n",
    "                visualizer.get_input_port(0))\n",
    "# end setup for visualization\n",
    "\n",
    "diagram = builder.Build()\n",
    "\n",
    "simulator = Simulator(diagram)\n",
    "#simulator.set_target_realtime_rate(1.0)\n",
    "context = simulator.get_mutable_context()\n",
    "\n",
    "duration = 4.0 if running_as_notebook else 0.1 # sets a shorter duration during testing\n",
    "visualizer.start_recording()\n",
    "print('simulating...')\n",
    "for i in range(5):\n",
    "    context.SetTime(0.)\n",
    "    context.SetContinuousState(np.random.randn(12,))\n",
    "    simulator.Initialize()\n",
    "    simulator.AdvanceTo(duration)\n",
    "visualizer.publish_recording()\n",
    "if running_as_notebook:\n",
    "    visualizer.vis.render_static()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
